import torch
import torch.nn as nn
import torch.optim as optim
import time
import cv2

from netCreater import *
from DataLoader import *
from render import *


def train(net, opt, data):
    # 设置输入
    inp = data[0].cuda()
    outp = data[1].cuda()
    # 预测和计算损失
    opt.zero_grad()
    y = net(inp)
    loss = criterion(y, outp)
    # 反向传播
    loss.backward()
    opt.step()
    return loss.item()


def test(net, data):
    with torch.no_grad():
        # 设置输入
        inp = data[0].cuda()
        outp = data[1].cuda()
        # 预测和计算损失
        y = net.forward(inp)
        loss = criterion(y, outp)
    return loss.item()


def TrainAndShow(
    net,
    datas,
    expectation=-1e60,
    learningRate=1e-3,
    batchSize=32,
    patience=1e5,
    saveInterval=5,
    epoch=100,
):
    print("训练开始.")
    # 训练参数
    if 1:
        opt = optim.AdamW(net.parameters(), lr=learningRate)
    else:
        opt = optim.SGD(net.parameters(), lr=learningRate)
    trTimes = roof(datas.trLength, batchSize)
    teTimes = roof(datas.teLength, batchSize)
    # 过程数据
    tStart = time.time()
    trainLoss = []
    testLoss = []
    # 提前终止机制
    badTime = 0
    loss = 1000
    lossprev = 1000
    # 可视化
    graph = np.zeros((480 * 2, 640, 3))
    img1 = np.zeros((480, 640, 3))
    img2 = np.zeros((480, 640, 3))
    col1 = [[0, 0, 255]]
    col2 = [[0, 255, 0]]
    k = 0
    cnt = 0
    with open("log.cvs", "w") as f:
        f.write("epochs\ttime\ttrainloss\ttestloss\n")
    while loss > expectation and badTime < patience and k != ord("q") and epoch != 0:
        if epoch > 0:
            epoch -= 1
        ### 训练
        loss = 0
        for i in range(trTimes):
            loss += train(net, opt, datas.getTrain(batchSize)) / trTimes / batchSize
        trainLoss += [[time.time() - tStart, loss]]
        ### 测试
        loss = 0
        for i in range(teTimes):
            loss += test(net, datas.getTest(batchSize)) / teTimes / batchSize
        testLoss += [[time.time() - tStart, loss]]
        # 显示和交互
        t = str(time.time() - tStart)[:5]
        print("\n次数=", cnt, ",时间=" + str(t)[:5])
        print(">>训练集误差 = " + str(trainLoss[-1]))
        print(">>测试集误差 = " + str(testLoss[-1]))
        with open("log.cvs", "a") as f:
            f.write(
                str(cnt)
                + "\t"
                + t
                + "\t"
                + str(trainLoss[-1][1])
                + "\t"
                + str(testLoss[-1][1])
                + "\n"
            )
        # img1*=0
        # img2*=0
        # img1=drawGraph(img1,[trainLoss[-30:]],col1)
        # img2=drawGraph(img2,[testLoss[-30:]],col2)
        # graph[0:480,0:640]=img1
        # graph[480:960,0:640]=img2
        # cv2.imshow("LossCurve",graph)
        # k=cv2.waitKey(1)
        # 耐心机制
        if len(testLoss) > 2 and testLoss[-1][1] > testLoss[-2][1]:
            badTime += 1
        else:
            badTime = 0
        # 交互
        cnt += 1
        loss = testLoss[-1][1]
        if cnt % saveInterval == 0 and lossprev > loss or k == ord("s"):
            torch.save(net, "./model.pt")
            print("模型保存")
            lossprev = loss

    print("训练结束.")


def roof(a, b):
    if (a % b) == 0:
        return a // b
    return a // b + 1


class CustomLoss(nn.Module):
    def __init__(self):
        super(CustomLoss, self).__init__()

    def forward(self, y, t):
        loss = -(t * torch.log(y) + (1 - t) * torch.log(1 - y))
        return torch.mean(loss)


if __name__ == "__main__":
    # 准备数据集
    path = r"D:/pytorchPro/Detection/dataset"
    dataSet = DataLoader(path, width=WIDTH, height=HEIGHT, channal=CHANNEL)

    # 损失函数
    criterion = nn.BCELoss()

    # 读取网络
    if 0:
        net = torch.load("model.pt")
    else:
        net = MyNet()
    net.cuda()
    TrainAndShow(net, dataSet, learningRate=1e-3, epoch=-1)
    cv2.destroyAllWindows()

    ####### 测试
    # 获取输入输出
    inp = dataSet.getTest(1)[0]
    outp = net(inp.cuda())
    # 把输入输出转化成 cv.Mat
    inp = (inp[0].cpu().numpy()[0].T + 1) * 255 / 2
    inp = np.array(inp, dtype=np.uint8)

    outp = outp.detach().cpu().numpy()[0].T * 255
    outp = np.array(outp, dtype=np.uint8)
    cv2.imshow("inp", inp)
    cv2.imshow("outp", outp)
    cv2.waitKey(0)
